{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f94c3f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _testing.local_redpanda_broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2fbc62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "import asyncio\n",
    "from pathlib import Path\n",
    "from contextlib import asynccontextmanager\n",
    "from tempfile import TemporaryDirectory\n",
    "from typing import *\n",
    "\n",
    "import asyncer\n",
    "import nest_asyncio\n",
    "\n",
    "from fastkafka._components._subprocess import terminate_asyncio_process\n",
    "from fastkafka._components.helpers import in_notebook\n",
    "from fastkafka._components.logger import get_logger\n",
    "from fastkafka._components.meta import delegates, export, patch\n",
    "from fastkafka._testing.apache_kafka_broker import get_free_port, run_and_match, _get_unique_local_brokers_to_start, _start_and_stop_brokers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de5d98a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "\n",
    "from fastkafka._aiokafka_imports import AIOKafkaConsumer, AIOKafkaProducer\n",
    "from fastkafka._components.logger import suppress_timestamps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8cd1883",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "if in_notebook():\n",
    "    from tqdm.notebook import tqdm\n",
    "else:\n",
    "    from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d872ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ec8c10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0291208d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def get_redpanda_docker_cmd(\n",
    "    listener_port: int = 9092,\n",
    "    tag: str = \"v23.1.2\",\n",
    "    seastar_core: int = 1,\n",
    "    memory: str = \"1G\",\n",
    "    mode: str = \"dev-container\",\n",
    "    default_log_level: str = \"debug\",\n",
    ") -> List[str]:\n",
    "    \"\"\"\n",
    "    Generates a Docker CLI command to start redpanda container\n",
    "\n",
    "    Args:\n",
    "        listener_port: Port on which the clients (producers and consumers) can connect\n",
    "        tag: Tag of Redpanda image to use to start container\n",
    "        seastar_core: Core(s) to use byt Seastar (the framework Redpanda uses under the hood)\n",
    "        memory: The amount of memory to make available to Redpanda\n",
    "        mode: Mode to use to load configuration properties in container\n",
    "        default_log_level: Log levels to use for Redpanda\n",
    "    \"\"\"\n",
    "    redpanda_docker_cmd = [\n",
    "        \"docker\",\n",
    "        \"run\",\n",
    "        \"--rm\",\n",
    "        \"--name\",\n",
    "        f\"redpanda_{listener_port}\",\n",
    "        \"-p\",\n",
    "        f\"{listener_port}:{listener_port}\",\n",
    "        f\"docker.redpanda.com/redpandadata/redpanda:{tag}\",\n",
    "        \"redpanda\",\n",
    "        \"start\",\n",
    "        \"--kafka-addr\",\n",
    "        f\"internal://0.0.0.0:9090,external://0.0.0.0:{listener_port}\",\n",
    "        \"--advertise-kafka-addr\",\n",
    "        f\"internal://localhost:9090,external://localhost:{listener_port}\",\n",
    "        \"--smp\",\n",
    "        str(seastar_core),\n",
    "        \"--memory\",\n",
    "        memory,\n",
    "        \"--mode\",\n",
    "        mode,\n",
    "        \"--default-log-level\",\n",
    "        default_log_level,\n",
    "    ]\n",
    "    return redpanda_docker_cmd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfddd7b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "expected = \"docker run --rm --name redpanda_9093 -p 9093:9093 docker.redpanda.com/redpandadata/redpanda:latest redpanda start --kafka-addr internal://0.0.0.0:9090,external://0.0.0.0:9093 --advertise-kafka-addr internal://localhost:9090,external://localhost:9093 --smp 1 --memory 2G --mode dev-container --default-log-level debug\"\n",
    "actual = get_redpanda_docker_cmd(listener_port=9093, tag=\"latest\", memory=\"2G\")\n",
    "assert \" \".join(actual) == expected, actual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dc00a80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@export(\"fastkafka.testing\")\n",
    "class LocalRedpandaBroker:\n",
    "    \"\"\"LocalRedpandaBroker class, used for running unique redpanda brokers in tests to prevent topic clashing.\"\"\"\n",
    "\n",
    "    @delegates(get_redpanda_docker_cmd, keep=True)\n",
    "    def __init__(\n",
    "        self,\n",
    "        topics: Iterable[str] = [],\n",
    "        *,\n",
    "        retries: int = 3,\n",
    "        apply_nest_asyncio: bool = False,\n",
    "        **kwargs: Dict[str, Any],\n",
    "    ):\n",
    "        \"\"\"Initialises the LocalRedpandaBroker object\n",
    "\n",
    "        Args:\n",
    "            topics: List of topics to create after sucessfull redpanda broker startup\n",
    "            retries: Number of retries to create redpanda service\n",
    "            apply_nest_asyncio: set to True if running in notebook\n",
    "            listener_port: Port on which the clients (producers and consumers) can connect\n",
    "            tag: Tag of Redpanda image to use to start container\n",
    "            seastar_core: Core(s) to use byt Seastar (the framework Redpanda uses under the hood)\n",
    "            memory: The amount of memory to make available to Redpanda\n",
    "            mode: Mode to use to load configuration properties in container\n",
    "            default_log_level: Log levels to use for Redpanda\n",
    "        \"\"\"\n",
    "        self.redpanda_kwargs = kwargs\n",
    "\n",
    "        if \"listener_port\" not in self.redpanda_kwargs:\n",
    "            self.redpanda_kwargs[\"listener_port\"] = 9092  # type: ignore\n",
    "\n",
    "        self.retries = retries\n",
    "        self.apply_nest_asyncio = apply_nest_asyncio\n",
    "        self.temporary_directory: Optional[TemporaryDirectory] = None\n",
    "        self.temporary_directory_path: Optional[Path] = None\n",
    "        self.redpanda_task: Optional[asyncio.subprocess.Process] = None\n",
    "        self._is_started = False\n",
    "        self.topics: Iterable[str] = topics\n",
    "\n",
    "    @property\n",
    "    def is_started(self) -> bool:\n",
    "        \"\"\"Property indicating whether the LocalRedpandaBroker object is started.\n",
    "\n",
    "        The is_started property indicates if the LocalRedpandaBroker object is currently \n",
    "        in a started state. This implies that Redpanda docker container has sucesfully \n",
    "        started and is ready for handling events.\n",
    "\n",
    "        Returns:\n",
    "            bool: True if the object is started, False otherwise.\n",
    "        \"\"\"\n",
    "        return self._is_started\n",
    "\n",
    "    @classmethod\n",
    "    async def _check_deps(cls) -> None:\n",
    "        \"\"\"Prepares the environment for running redpanda brokers.\n",
    "        Returns:\n",
    "           None\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _start(self) -> str:\n",
    "        \"\"\"Starts a local redpanda broker instance asynchronously\n",
    "        Returns:\n",
    "           Redpanda broker bootstrap server address in string format: add:port\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def start(self) -> str:\n",
    "        \"\"\"Starts a local redpanda broker instance synchronously\n",
    "        Returns:\n",
    "           Redpanda broker bootstrap server address in string format: add:port\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def stop(self) -> None:\n",
    "        \"\"\"Stops a local redpanda broker instance synchronously\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _stop(self) -> None:\n",
    "        \"\"\"Stops a local redpanda broker instance synchronously\n",
    "        Returns:\n",
    "           None\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def get_service_config_string(self, service: str, *, data_dir: Path) -> str:\n",
    "        \"\"\"Generates a configuration for a service\n",
    "        Args:\n",
    "            data_dir: Path to the directory where the zookeepeer instance will save data\n",
    "            service: \"redpanda\", defines which service to get config string for\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _start_redpanda(self) -> None:\n",
    "        \"\"\"Start a local redpanda broker\n",
    "        Returns:\n",
    "           None\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    async def _create_topics(self) -> None:\n",
    "        \"\"\"Create missing topics in local redpanda broker\n",
    "        Returns:\n",
    "           None\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __enter__(self) -> str:\n",
    "        return self.start()\n",
    "\n",
    "    def __exit__(self, *args: Any, **kwargs: Any) -> None:\n",
    "        self.stop()\n",
    "\n",
    "    async def __aenter__(self) -> str:\n",
    "        return await self._start()\n",
    "\n",
    "    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:\n",
    "        await self._stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04f2dc8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<fastkafka.testing.LocalRedpandaBroker>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "broker = LocalRedpandaBroker()\n",
    "broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94f530a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "async def check_docker(tag: str = \"v23.1.2\") -> bool:\n",
    "    \"\"\"\n",
    "    Checks if a Docker image with the specified tag is available.\n",
    "\n",
    "    Args:\n",
    "        tag: The tag of the Docker image to check. Defaults to \"v23.1.2\".\n",
    "\n",
    "    Returns:\n",
    "        bool: True if the Docker image is available; False otherwise.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        docker_task = await run_and_match(\n",
    "            \"docker\",\n",
    "            \"pull\",\n",
    "            f\"docker.redpanda.com/redpandadata/redpanda:{tag}\",\n",
    "            pattern=f\"docker.redpanda.com/redpandadata/redpanda:{tag}\",\n",
    "        )\n",
    "        return True\n",
    "    except Exception as e:\n",
    "        logger.debug(f\"Error in check_docker() : {e}\")\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecc3b8df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await check_docker()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f6673e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch(cls_method=True)  # type: ignore\n",
    "async def _check_deps(cls: LocalRedpandaBroker) -> None:\n",
    "    if not await check_docker():\n",
    "        raise RuntimeError(\n",
    "            \"Docker installation not found! Please install docker manually and retry.\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c20e8f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "broker = LocalRedpandaBroker()\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    await broker._check_deps()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80275185",
   "metadata": {},
   "outputs": [],
   "source": [
    "# docker run --rm --name redpanda_9092 -p 9092:9092 docker.redpanda.com/redpandadata/redpanda 'redpanda start --kafka-addr internal://0.0.0.0:9090,external://0.0.0.0:9092 --advertise-kafka-addr internal://localhost:9090,external://localhost:9092'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae79c68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _start_redpanda(self: LocalRedpandaBroker, service: str = \"redpanda\") -> None:\n",
    "    logger.info(f\"Starting {service}...\")\n",
    "\n",
    "    if self.temporary_directory_path is None:\n",
    "        raise ValueError(\n",
    "            \"LocalRedpandaBroker._start_redpanda(): self.temporary_directory_path is None, did you initialise it?\"\n",
    "        )\n",
    "\n",
    "    configs_tried: List[Dict[str, Any]] = []\n",
    "\n",
    "    for i in range(self.retries + 1):\n",
    "        configs_tried = configs_tried + [getattr(self, f\"{service}_kwargs\").copy()]\n",
    "\n",
    "        redpanda_docker_cmd = get_redpanda_docker_cmd(**self.redpanda_kwargs)  # type: ignore\n",
    "\n",
    "        try:\n",
    "            service_task = await run_and_match(\n",
    "                *redpanda_docker_cmd,\n",
    "                capture=\"stderr\",\n",
    "                pattern=\"Bootstrap complete\",\n",
    "                timeout=30,\n",
    "            )\n",
    "        except Exception as e:\n",
    "            logger.info(\n",
    "                f\"{service} startup failed, generating a new port and retrying...\"\n",
    "            )\n",
    "            port = get_free_port()\n",
    "            self.redpanda_kwargs[\"listener_port\"] = port  # type: ignore\n",
    "\n",
    "            logger.info(f\"port={port}\")\n",
    "        else:\n",
    "            setattr(self, f\"{service}_task\", service_task)\n",
    "            return\n",
    "\n",
    "    raise ValueError(f\"Could not start {service} with params: {configs_tried}\")\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _create_topics(self: LocalRedpandaBroker) -> None:\n",
    "    listener_port = self.redpanda_kwargs.get(\"listener_port\", 9092)\n",
    "\n",
    "    async with asyncer.create_task_group() as tg:\n",
    "        processes = [\n",
    "            tg.soonify(run_and_match)(\n",
    "                \"docker\",\n",
    "                \"exec\",\n",
    "                f\"redpanda_{listener_port}\",\n",
    "                \"rpk\",\n",
    "                \"topic\",\n",
    "                \"create\",\n",
    "                topic,\n",
    "                pattern=topic,\n",
    "                timeout=10,\n",
    "            )\n",
    "            for topic in self.topics\n",
    "        ]\n",
    "\n",
    "    try:\n",
    "        return_values = [\n",
    "            await asyncio.wait_for(process.value.wait(), 30) for process in processes\n",
    "        ]\n",
    "        if any(return_value != 0 for return_value in return_values):\n",
    "            raise ValueError(\"Could not create missing topics!\")\n",
    "    except asyncio.TimeoutError as _:\n",
    "        raise ValueError(\"Timed out while creating missing topics!\")\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _start(self: LocalRedpandaBroker) -> str:\n",
    "    await self._check_deps()\n",
    "\n",
    "    self.temporary_directory = TemporaryDirectory()\n",
    "    self.temporary_directory_path = Path(self.temporary_directory.__enter__())\n",
    "\n",
    "    await self._start_redpanda()\n",
    "    await asyncio.sleep(5)\n",
    "\n",
    "    listener_port = self.redpanda_kwargs.get(\"listener_port\", 9092)\n",
    "    bootstrap_server = f\"127.0.0.1:{listener_port}\"\n",
    "    logger.info(f\"Local Redpanda broker up and running on {bootstrap_server}\")\n",
    "\n",
    "    await self._create_topics()\n",
    "\n",
    "    self._is_started = True\n",
    "\n",
    "    return bootstrap_server\n",
    "\n",
    "\n",
    "@patch\n",
    "async def _stop(self: LocalRedpandaBroker) -> None:\n",
    "    logger.info(f\"Stopping redpanda...\")\n",
    "    await terminate_asyncio_process(self.redpanda_task)  # type: ignore\n",
    "    logger.info(f\"Redpanda stopped.\")\n",
    "    self.temporary_directory.__exit__(None, None, None)  # type: ignore\n",
    "    self._is_started = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f09cf87",
   "metadata": {},
   "outputs": [],
   "source": [
    "broker = LocalRedpandaBroker()\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    async with broker:\n",
    "        pass\n",
    "\n",
    "    print(\"*\" * 50 + \"REDPANDA LOGS\" + \"+\" * 50)\n",
    "    redpanda_output, _ = await broker.redpanda_task.communicate()\n",
    "    print(redpanda_output.decode(\"UTF-8\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b98298a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def start(self: LocalRedpandaBroker) -> str:\n",
    "    \"\"\"Starts a local redpanda broker instance synchronously\n",
    "    Returns:\n",
    "       Redpanda broker bootstrap server address in string format: add:port\n",
    "    \"\"\"\n",
    "    logger.info(f\"{self.__class__.__name__}.start(): entering...\")\n",
    "    try:\n",
    "        # get or create loop\n",
    "        try:\n",
    "            loop = asyncio.get_event_loop()\n",
    "        except RuntimeError as e:\n",
    "            logger.warning(\n",
    "                f\"{self.__class__.__name__}.start(): RuntimeError raised when calling asyncio.get_event_loop(): {e}\"\n",
    "            )\n",
    "            logger.warning(\n",
    "                f\"{self.__class__.__name__}.start(): asyncio.new_event_loop()\"\n",
    "            )\n",
    "            loop = asyncio.new_event_loop()\n",
    "\n",
    "        # start redpanda broker in the loop\n",
    "\n",
    "        if loop.is_running():\n",
    "            if self.apply_nest_asyncio:\n",
    "                logger.warning(\n",
    "                    f\"{self.__class__.__name__}.start(): ({loop}) is already running!\"\n",
    "                )\n",
    "                logger.warning(\n",
    "                    f\"{self.__class__.__name__}.start(): calling nest_asyncio.apply()\"\n",
    "                )\n",
    "                nest_asyncio.apply(loop)\n",
    "            else:\n",
    "                msg = f\"{self.__class__.__name__}.start(): ({loop}) is already running! Use 'apply_nest_asyncio=True' when creating 'LocalRedpandaBroker' to prevent this.\"\n",
    "                logger.error(msg)\n",
    "                raise RuntimeError(msg)\n",
    "\n",
    "        retval = loop.run_until_complete(self._start())\n",
    "        logger.info(f\"{self.__class__}.start(): returning {retval}\")\n",
    "        return retval\n",
    "    finally:\n",
    "        logger.info(f\"{self.__class__.__name__}.start(): exited.\")\n",
    "\n",
    "\n",
    "@patch\n",
    "def stop(self: LocalRedpandaBroker) -> None:\n",
    "    \"\"\"Stops a local redpanda broker instance synchronously\"\"\"\n",
    "    logger.info(f\"{self.__class__.__name__}.stop(): entering...\")\n",
    "    try:\n",
    "        if not self._is_started:\n",
    "            raise RuntimeError(\n",
    "                \"LocalRedpandaBroker not started yet, please call LocalRedpandaBroker.start() before!\"\n",
    "            )\n",
    "\n",
    "        loop = asyncio.get_event_loop()\n",
    "        loop.run_until_complete(self._stop())\n",
    "    finally:\n",
    "        logger.info(f\"{self.__class__.__name__}.stop(): exited.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c92191",
   "metadata": {},
   "outputs": [],
   "source": [
    "broker = LocalRedpandaBroker(apply_nest_asyncio=True)\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    with broker:\n",
    "        print(\"Hello world!\")\n",
    "\n",
    "    print(\"*\" * 50 + \"REDPANDA LOGS\" + \"+\" * 50)\n",
    "    redpanda_output, _ = await broker.redpanda_task.communicate()\n",
    "    print(redpanda_output.decode(\"UTF-8\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94867ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_topic = \"test-topic\"\n",
    "test_msg = b\"test-msg\"\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    with LocalRedpandaBroker(\n",
    "        topics=[test_topic], apply_nest_asyncio=True\n",
    "    ) as bootstrap_server:\n",
    "        consumer = AIOKafkaConsumer(test_topic, bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        producer = AIOKafkaProducer(bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        await consumer.start()\n",
    "        await producer.start()\n",
    "\n",
    "        try:\n",
    "            await producer.send_and_wait(test_topic, test_msg)\n",
    "            msg = await consumer.getone()\n",
    "            assert msg, value == test_msg\n",
    "        finally:\n",
    "            await consumer.stop()\n",
    "            await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68f7183",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_topic = \"test-topic\"\n",
    "test_msg = b\"test-msg\"\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    async with LocalRedpandaBroker(topics=[test_topic]) as bootstrap_server:\n",
    "        consumer = AIOKafkaConsumer(test_topic, bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        producer = AIOKafkaProducer(bootstrap_servers=bootstrap_server)\n",
    "\n",
    "        await consumer.start()\n",
    "        await producer.start()\n",
    "\n",
    "        try:\n",
    "            await producer.send_and_wait(test_topic, test_msg)\n",
    "            msg = await consumer.getone()\n",
    "            assert msg, value == test_msg\n",
    "        finally:\n",
    "            await consumer.stop()\n",
    "            await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6335b063",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def start_redpanda_brokers(\n",
    "    kafka_brokers_name: str,\n",
    "    kafka_brokers: List[Dict[str, Dict[str, Any]]],\n",
    "    duplicate_ok: bool = False,\n",
    "    ignore_nonlocal_brokers: bool = False,\n",
    ") -> AsyncIterator[None]:\n",
    "    unique_local_brokers_to_start = await _get_unique_local_brokers_to_start(\n",
    "        kafka_brokers_name=kafka_brokers_name,\n",
    "        kafka_brokers=kafka_brokers,\n",
    "        duplicate_ok=duplicate_ok,\n",
    "        ignore_nonlocal_brokers=ignore_nonlocal_brokers,\n",
    "    )\n",
    "\n",
    "    brokers = [\n",
    "        LocalRedpandaBroker(listener_port=broker[1])  # type: ignore\n",
    "        for broker in unique_local_brokers_to_start\n",
    "    ]\n",
    "\n",
    "    async with _start_and_stop_brokers(brokers=brokers):\n",
    "        yield"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "858c8a09",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_broker_1 = dict(\n",
    "    development=dict(url=\"localhost\", port=39092),\n",
    "    production=dict(url=\"prod.server_1\", port=9092),\n",
    ")\n",
    "\n",
    "kafka_broker_2 = dict(\n",
    "    development=dict(url=\"localhost\", port=39093),\n",
    "    production=dict(url=\"prod.server_1\", port=9092),\n",
    ")\n",
    "\n",
    "kafka_brokers = [kafka_broker_1, kafka_broker_2]\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "    async with start_redpanda_brokers(\n",
    "        kafka_brokers_name=\"development\",\n",
    "        kafka_brokers=kafka_brokers,\n",
    "    ):\n",
    "        print(\"Succesfully started multiple brokers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc7ebbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_brokers = [\n",
    "    dict(\n",
    "        development=dict(url=\"localhost\", port=39092),\n",
    "    ),\n",
    "    dict(\n",
    "        development=dict(url=\"localhost\", port=39092),\n",
    "    ),\n",
    "]\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "\n",
    "    with pytest.raises(ValueError) as e:\n",
    "        async with start_redpanda_brokers(\n",
    "            kafka_brokers_name=\"development\",\n",
    "            kafka_brokers=kafka_brokers,\n",
    "            duplicate_ok=False,\n",
    "        ):\n",
    "            print(\"If this is printed then the test didn't pass\")\n",
    "    print(e)\n",
    "    assert \"Duplicate kafka_brokers are found\" in str(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "979b4406",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_brokers = [\n",
    "    dict(\n",
    "        development=dict(url=\"127.0.0.1\", port=39092),\n",
    "    ),\n",
    "    dict(\n",
    "        development=dict(url=\"some_other_host\", port=39092),\n",
    "    ),\n",
    "]\n",
    "\n",
    "# Run test case only if docker is installed\n",
    "if await check_docker():\n",
    "\n",
    "    with pytest.raises(ValueError) as e:\n",
    "        async with start_redpanda_brokers(\n",
    "            kafka_brokers_name=\"development\",\n",
    "            kafka_brokers=kafka_brokers,\n",
    "            ignore_nonlocal_brokers=False,\n",
    "        ):\n",
    "            print(\"If this is printed then the test didn't pass\")\n",
    "    print(e)\n",
    "    assert \"URL values other than 0.0.0.0, 127.0.0.1, localhost are found\" in str(e)\n",
    "\n",
    "    async with start_redpanda_brokers(\n",
    "        kafka_brokers_name=\"development\",\n",
    "        kafka_brokers=kafka_brokers,\n",
    "        ignore_nonlocal_brokers=True,\n",
    "    ):\n",
    "        print(\"Starting again with 'ignore_nonlocal_brokers=True'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6404c272",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
